from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any, Union

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message] = Field(default_factory=list)
    query: Optional[str] = None  # 可选的RAG查询
    audio_output: bool = False   # 是否需要音频输出

class ChatResponse(BaseModel):
    content: str
    role: str = "assistant"
    content_type: str = "text"  # 内容类型：text或audio
    audio_url: Optional[str] = None  # 音频URL，当content_type为audio时提供
    
class StreamChunk(BaseModel):
    content: str
    content_type: str = "text"  # 内容类型：text或audio
    audio_url: Optional[str] = None  # 音频URL，当content_type为audio时提供
    finish_reason: Optional[str] = None

class WebSocketMessage(BaseModel):
    """WebSocket消息基础模型"""
    type: str  # 消息类型：chat, ping, clear_session等
    data: Optional[Dict[str, Any]] = None  # 消息数据
    timestamp: Optional[str] = None  # 时间戳

class WebSocketResponse(BaseModel):
    """WebSocket响应模型"""
    type: str  # 响应类型：session_info, response_start, response_chunk, response_end, tts_start, tts_end, tts_error, pong, error等
    content: Optional[str] = None  # 响应内容
    content_type: Optional[str] = "text"  # 内容类型：text或audio
    audio_url: Optional[str] = None  # 音频文件URL
    session_id: Optional[str] = None  # 会话ID
    finish_reason: Optional[str] = None  # 完成原因
    message: Optional[str] = None  # 状态或错误消息
    timestamp: Optional[str] = None  # 时间戳